from convdiffsolverbase import ConvDiffSolverBase
from dolfin import *
import numpy

class Crosswind(ConvDiffSolverBase):
    '''Original crosswind method by Johnson(1987).'''
    def __init__(self):
        self.name = 'Crosswind'
        ConvDiffSolverBase.__init__(self)
    
    def solve(self,problem,get_tau,get_sigma):
        '''SUPG stab. parameter tau, and crosswind stab parameters are obtained by calling appropriate functions.'''
        
        problemName, mesh, eps, BCValues, b, f, BCIndicators = problem.get_vars()
        
        V = FunctionSpace(mesh,'CG',1) 
        VV = VectorFunctionSpace(mesh,'CG',1)
        DGV = FunctionSpace(mesh,'DG',0)
        
        u = TrialFunction(V)
        varphi = TestFunction(V)
        
        v = interpolate(b,VV)

        tau = Function(DGV)
        tau.vector()[:] = get_tau(v,eps) # should be the same as in upwind diffusion
        
        sigma = Function(DGV)
        sigma.vector()[:] = get_sigma(v,eps)

        pert = conditional(
                           gt(sqrt(inner(v,v)),1E-16),
                           Constant(1.0),
                           Constant(0.) 
                          ) # a way to realize swith on |v| == 0.
                        
        pMatrix = conditional(
                             gt(sqrt(inner(v,v)),1E-16),
                             Identity(2)-outer(v,v)/inner(v,v),
                             Constant(0.)*Identity(2),
                             ) # I - v X v
        
        a = eps*inner(nabla_grad(u),nabla_grad(varphi))*dx + inner(v,nabla_grad(u))*varphi*dx +\
            tau*pert*inner(dot(nabla_grad(u),outer(v,v)),nabla_grad(varphi))*dx+\
            sigma*inner(dot(nabla_grad(u),pMatrix),nabla_grad(varphi))*dx       

        l = f*varphi*dx + tau*pert*f*inner(v,nabla_grad(varphi))*dx
    
        bcs = [DirichletBC(V,value,where) for value, where in zip(BCValues, BCIndicators)]

        u = Function(V)
        A,L = assemble_system(a,l,bcs)
        solve(A,u.vector(),L)
            
        plot(u,interactive=True)
        self.save(u,problem)
